#include <sys/resource.h>
#include <stdlib.h>
#include <stdio.h>

void print_rlimit()
{
	struct rlimit limit;
	if (getrlimit(RLIMIT_AS, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_AS rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);
	if (getrlimit(RLIMIT_CORE, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_CORE rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);
	if (getrlimit(RLIMIT_CPU, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_CPU rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);
	if (getrlimit(RLIMIT_DATA, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_DATA rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);
	if (getrlimit(RLIMIT_FSIZE, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_FSIZE rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);
	if (getrlimit(RLIMIT_LOCKS, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_LOCKS rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);
	if (getrlimit(RLIMIT_MEMLOCK, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_MEMLOCK rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);
	if (getrlimit(RLIMIT_NOFILE, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_NOFILE rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);
	if (getrlimit(RLIMIT_NPROC, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_NPROC rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);
	if (getrlimit(RLIMIT_RSS, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_RSS rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);

	if (getrlimit(RLIMIT_STACK, &limit) < 0) {
		perror("getrlimit error");
		exit(0);
	}
	printf("RLIMIT_STACK rlim_cur:%ld rlim_max:%ld \n", limit.rlim_cur,
	       limit.rlim_max);

}

void set_rlimit(long time_limit, long memory_limit)
{

	struct rlimit limit;

	limit.rlim_cur = time_limit;
	limit.rlim_max = time_limit;
	if (setrlimit(RLIMIT_CPU, &limit))
		perror("failed to set RLIMIT_CPU");

	//TODO
/*	
	limit.rlim_cur = memory_limit;
	limit.rlim_max = memory_limit;
	setrlimit(RLIMIT_AS, &limit);
	// SIGKILL -> SIGENV	
*/

	limit.rlim_cur = memory_limit;
	limit.rlim_max = memory_limit;
	if (setrlimit(RLIMIT_DATA, &limit))
		perror("failed to set RLIMIT_DATA");

	limit.rlim_cur = memory_limit;
	limit.rlim_max = memory_limit;
	setrlimit(RLIMIT_MEMLOCK, &limit);

	//no open more files except stds
	limit.rlim_cur = 4;
	limit.rlim_max = 4;
	setrlimit(RLIMIT_NOFILE, &limit);

	// no core file dump
	limit.rlim_cur = 0;
	limit.rlim_max = 0;
	setrlimit(RLIMIT_CORE, &limit);

	// no child fork
	limit.rlim_cur = 1;
	limit.rlim_max = 1;
	setrlimit(RLIMIT_NPROC, &limit);

	limit.rlim_cur = 1024;
	limit.rlim_max = 1024;
	setrlimit(RLIMIT_FSIZE, &limit);

	limit.rlim_cur = -1;
	limit.rlim_max = -1;
	setrlimit(RLIMIT_LOCKS, &limit);

	limit.rlim_cur = memory_limit;
	limit.rlim_max = memory_limit;
	setrlimit(RLIMIT_RSS, &limit);
	limit.rlim_cur = memory_limit;
	limit.rlim_max = memory_limit;
	if (setrlimit(RLIMIT_STACK, &limit) < 0) {
		perror("Can not setrlimit RLIMIT_STACK");
	}
}
